package com.zst.sw.server.network;

import com.zst.sw.server.pipeline.TraceSegmentDataPipeline;
import io.grpc.stub.StreamObserver;
import lombok.extern.slf4j.Slf4j;
import org.apache.skywalking.apm.network.common.v3.Commands;
import org.apache.skywalking.apm.network.language.agent.v3.SegmentCollection;
import org.apache.skywalking.apm.network.language.agent.v3.SegmentObject;
import org.apache.skywalking.apm.network.language.agent.v3.TraceSegmentReportServiceGrpc;

import java.util.List;

@Slf4j
public class TraceSegmentReportServiceHandler extends TraceSegmentReportServiceGrpc.TraceSegmentReportServiceImplBase {
    private TraceSegmentDataPipeline pipeline;

    public TraceSegmentReportServiceHandler(TraceSegmentDataPipeline pipeline) {
        this.pipeline = pipeline;
    }

    @Override
    public StreamObserver<SegmentObject> collect(StreamObserver<Commands> responseObserver) {
        return new StreamObserver<SegmentObject>() {
            @Override
            public void onNext(SegmentObject segmentObject) {
                pipeline.publish(segmentObject);
            }

            @Override
            public void onError(Throwable throwable) {
                log.error("trace segment report collect error", throwable);
            }

            @Override
            public void onCompleted() {
                responseObserver.onNext(Commands.newBuilder().build());
                responseObserver.onCompleted();
            }
        };
    }

    @Override
    public void collectInSync(SegmentCollection request, StreamObserver<Commands> responseObserver) {
        request.getSegmentsList().forEach(segmentObject -> pipeline.publish(segmentObject));

        responseObserver.onNext(Commands.newBuilder().build());
        responseObserver.onCompleted();
    }
}
